#include "Compress.h"
#include "Lzma2Dec.h"
#include "Lzma2Enc.h"
#include <windows.h>

namespace {

const int ChunkSizeMax = 4194304;
const int BufferSize = 65536;

struct InputStream
{
    ISeqInStream Input;
    HANDLE File;
    InputStream(const char* aName);
    ~InputStream();
};

struct OutputStream
{
    ISeqOutStream Output;
    HANDLE File;
    OutputStream(const char* aName);
    ~OutputStream();
};

SRes ReadFromStream(void* aStream, void* aBuffer, size_t* aSize)
{
    InputStream* Stream = reinterpret_cast<InputStream*>(aStream);
    size_t Requested = *aSize;

    *aSize = 0;

    do {
        DWORD Current = (Requested > ChunkSizeMax) ? ChunkSizeMax : (DWORD)Requested;
        DWORD Processed = 0;
        if (ReadFile(Stream->File, aBuffer, Current, &Processed, NULL)) {
            aBuffer = (void*)((Byte*)aBuffer + Processed);
            Requested -= Processed;
            *aSize += Processed;
        }
        else {
            return GetLastError();
        }
        if (Processed == 0) {
            Requested = 0;
        }
    } while (Requested > 0);

    return 0;
}

SRes ReadMinimumFromStream(ISeqInStream* aStream, void* aBuffer, size_t aSize)
{
    while (aSize != 0)
    {
        size_t Processed = aSize;
        RINOK(aStream->Read(aStream, aBuffer, &Processed));
        if (Processed == 0) {
            return SZ_ERROR_INPUT_EOF;
        }
        aBuffer = (void*)((Byte*)aBuffer + Processed);
        aSize -= Processed;
    }

    return SZ_OK;
}

size_t WriteToStream(void* aStream, const void* aBuffer, size_t aSize)
{
    OutputStream* Stream = reinterpret_cast<OutputStream*>(aStream);
    size_t Requested = aSize;

    aSize = 0;

    do {
        DWORD Current = (Requested > ChunkSizeMax) ? ChunkSizeMax : (DWORD)Requested;
        DWORD Processed = 0;
        if (WriteFile(Stream->File, aBuffer, Current, &Processed, NULL)) {
            aBuffer = (void*)((Byte*)aBuffer + Processed);
            Requested -= Processed;
            aSize += Processed;
        }
        else {
            return GetLastError();
        }
        if (Processed == 0) {
            Requested = 0;
        }
    } while (Requested > 0);

    return aSize;
}

InputStream::InputStream(const char* aName)
{
    File = CreateFile(aName, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
    Input.Read = File != INVALID_HANDLE_VALUE ? ReadFromStream : 0;
}

InputStream::~InputStream()
{
    CloseHandle(File);
}

OutputStream::OutputStream(const char* aName)
{
    File = CreateFile(aName, GENERIC_WRITE, FILE_SHARE_READ, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);
    Output.Write = File != INVALID_HANDLE_VALUE ? WriteToStream : 0;
}

OutputStream::~OutputStream()
{
    CloseHandle(File);
}

void* Malloc(void*, size_t aSize)
{
    return aSize == 0 ? static_cast<void*>(0) : malloc(aSize);
}

void Free(void*, void* aAddress)
{
    if (aAddress) {
        free(aAddress);
    }
}

ISzAlloc Alloc = { Malloc, Free };

SRes Encode(ISeqOutStream* aOutStream, ISeqInStream* aInStream, UInt64 aFileSize)
{
    CLzmaEncHandle Encoder = LzmaEnc_Create(&Alloc);

    if (Encoder == 0) {
        return SZ_ERROR_MEM;
    }

    CLzmaEncProps Properties;
    LzmaEncProps_Init(&Properties);

    SRes Result = LzmaEnc_SetProps(Encoder, &Properties);

    if (Result == SZ_OK) {
        Byte Header[LZMA_PROPS_SIZE + 8];
        size_t HeaderSize = LZMA_PROPS_SIZE;

        Result = LzmaEnc_WriteProperties(Encoder, Header, &HeaderSize);

        if (Result == SZ_OK) {
            for (int Index = 0; Index < 8; ++Index) {
                Header[HeaderSize++] = (Byte)(aFileSize >> (8 * Index));
            }
            if (aOutStream->Write(aOutStream, Header, HeaderSize) != HeaderSize) {
                Result = SZ_ERROR_WRITE;
            }
            else {
                Result = LzmaEnc_Encode(Encoder, aOutStream, aInStream, 0, &Alloc, &Alloc);
            }
        }
    }

    LzmaEnc_Destroy(Encoder, &Alloc, &Alloc);
    return Result;
}

SRes DecodeWork(CLzmaDec* aState, ISeqOutStream* aOutStream, ISeqInStream* aInStream, UInt64 aUnpackSize)
{
    int ThereIsSize = (aUnpackSize != (UInt64)(Int64)-1);
    Byte InputBuffer[BufferSize];
    Byte OutputBuffer[BufferSize];
    size_t InputPosition = 0, InputSize = 0, OutputPosition = 0;
    LzmaDec_Init(aState);

    for (;;)
    {
        if (InputPosition == InputSize) {
            InputSize = BufferSize;
            RINOK(aInStream->Read(aInStream, InputBuffer, &InputSize));
            InputPosition = 0;
        }
        {
        SRes Result;
        SizeT InputProcessed = InputSize - InputPosition;
        SizeT OutputProcessed = BufferSize - OutputPosition;
        ELzmaFinishMode FinishMode = LZMA_FINISH_ANY;
        ELzmaStatus Status;
        if (ThereIsSize && OutputProcessed > aUnpackSize) {
            OutputProcessed = (SizeT)aUnpackSize;
            FinishMode = LZMA_FINISH_END;
        }

        Result = LzmaDec_DecodeToBuf(aState, OutputBuffer + OutputPosition, &OutputProcessed, InputBuffer + InputPosition, &InputProcessed, FinishMode, &Status);
        InputPosition += InputProcessed;
        OutputPosition += OutputProcessed;
        aUnpackSize -= OutputProcessed;

        if (aOutStream) {
            if (aOutStream->Write(aOutStream, OutputBuffer, OutputPosition) != OutputPosition) {
                return SZ_ERROR_WRITE;
            }
        }

        OutputPosition = 0;

        if (Result != SZ_OK || ThereIsSize && aUnpackSize == 0)
            return Result;

        if (InputProcessed == 0 && OutputProcessed == 0) {
            if (ThereIsSize || Status != LZMA_STATUS_FINISHED_WITH_MARK) {
                return SZ_ERROR_DATA;
            }
            return Result;
        }
    }
    }
}

SRes Decode(ISeqOutStream* aOutStream, ISeqInStream* aInStream)
{
    unsigned char Header[LZMA_PROPS_SIZE + 8];

    RINOK(ReadMinimumFromStream(aInStream, Header, sizeof(Header)));

    UInt64 UnpackSize = 0;
    for (int Index = 0; Index < 8; ++Index)
        UnpackSize += (UInt64)Header[LZMA_PROPS_SIZE + Index] << (Index * 8);

    CLzmaDec State;
    LzmaDec_Construct(&State);
    RINOK(LzmaDec_Allocate(&State, Header, LZMA_PROPS_SIZE, &Alloc));
    SRes Result = DecodeWork(&State, aOutStream, aInStream, UnpackSize);
    LzmaDec_Free(&State, &Alloc);

    return Result;
}

} // anonymous namespace

int EncodeFile(const char* aInputName, const char* aOutputName)
{
    InputStream Input(aInputName);
    if (Input.File == INVALID_HANDLE_VALUE) {
        return GetLastError();
    }

    OutputStream Output(aOutputName);
    if (Output.File == INVALID_HANDLE_VALUE) {
        return GetLastError();
    }

    DWORD High;
    DWORD Low = GetFileSize(Input.File, &High);
    if (Low == 0xFFFFFFFF)
    {
        DWORD Check = GetLastError();
        if (Check != NO_ERROR)
            return Check;
    }

    UInt64 Size = (((UInt64)High) << 32) + Low;

    return Encode(&Output.Output, &Input.Input, Size);
}

int DecodeFile(const char* aInputName, const char* aOutputName)
{
    InputStream Input(aInputName);
    if (Input.File == INVALID_HANDLE_VALUE) {
        return GetLastError();
    }

    OutputStream Output(aOutputName);
    if (Output.File == INVALID_HANDLE_VALUE) {
        return GetLastError();
    }

    return Decode(&Output.Output, &Input.Input);
}
